import numpy as np
from math import sin, cos, radians, sqrt
import matplotlib.pyplot as plt

# --- I. 物理常数和弹丸参数 ---
# 国际标准大气模型参数
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 温度递减率 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 气体常数 J/(kg*K) 
P_TERM = (G / (R * L)) - 1 # 密度公式指数项

# 弹丸参数 (假设值，与之前保持一致)
C_SOUND = 338.0      
M = 5.0              
A = 0.0095           

# Cd-Ma 关系数据 (来自提供的曲线图)
MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) 
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) 

# --- II. 动态物理模型函数 ---

def get_drag_coefficient(v, c_sound):
    """根据马赫数线性插值计算 Cd。"""
    mach_number = v / c_sound
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS) # 考虑速度对Cd的关系
    return Cd

def get_air_density(z, model_type):
    """
    根据模型类型计算空气密度。
    model_type = 'dynamic' (R1): 考虑 ρ(z) 变化
    model_type = 'static' (R0): 使用海平面密度 ρ0
    """
    if model_type == 'static':
        # 模型 0 (R0): 不考虑海拔对rho的影响，使用海平面密度
        return RHO0
    
    elif model_type == 'dynamic':
        # 模型 1 (R1): 考虑海拔对rho的影响
        h = max(0.0, z) 
        temp_ratio = (T0 - L * h) / T0
        if temp_ratio <= 0:
            return 0.0 
        rho = RHO0 * (temp_ratio) ** P_TERM # 国际标准大气模型
        return rho
    else:
        # 如果模型类型输入有误，则抛出异常
        raise ValueError("Invalid model_type. Must be 'dynamic' or 'static'.")

# --- III. 弹道微分方程 (d S/d t) ---

def derivatives(S, params, model_type):
    """
    计算状态向量 S = [x, z, vx, vz] 对时间 t 的导数 dS/dt。
    """
    x, z, vx, vz = S
    m, A, g, c_sound = params

    v = sqrt(vx**2 + vz**2) 
    
    rho = get_air_density(z, model_type) 
    
    # Cd 动态依赖于速度
    Cd = get_drag_coefficient(v, c_sound) 
    
    drag_coeff = -(Cd * rho * A) / (2.0 * m) 
    
    dvx_dt = drag_coeff * v * vx 
    dvz_dt = drag_coeff * v * vz - g 
    
    return np.array([vx, vz, dvx_dt, dvz_dt])

# --- IV. RK4 积分函数 ---

def rk4_step(S, dt, params, model_type):
    """执行一步 RK4 积分，并传入 model_type。"""
    
    # K1 = h * f(t, S)
    K1 = dt * derivatives(S, params, model_type)
    
    # K2 = h * f(t + h/2, S + K1/2)
    K2 = dt * derivatives(S + K1 / 2.0, params, model_type)
    
    # K3 = h * f(t + h/2, S + K2/2)
    K3 = dt * derivatives(S + K2 / 2.0, params, model_type)
    
    # K4 = h * f(t + h, S + K3)
    K4 = dt * derivatives(S + K3, params, model_type)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

# --- V. 弹道模拟核心函数 (返回射程 R) ---

def calculate_range(v0, theta_deg, model_type, dt=0.01):
    """
    使用 RK4 算法模拟弹道轨迹，并返回射程 R。
    """
    params = (M, A, G, C_SOUND)
    theta_rad = radians(theta_deg)
    
    S = np.array([
        0.0, 0.0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    # 积分循环
    while S[1] >= -0.01: 
        S_new = rk4_step(S, dt, params, model_type)
        
        # 检查是否跨越地面
        if S_new[1] < 0 and S[1] >= 0:
            # 停止在落地点附近
            break 
            
        S = S_new

    R_final = S[0] # 返回 x 坐标作为射程
    return R_final

# --- VI. 交互式运行和结果计算 ---

def interactive_sensitivity_analysis(dt=0.01):
    
    print("\n--- 🎯 海拔对射程影响的敏感度分析 ---")
    print("请手动输入射击参数，程序将计算 R1 (动态rho) 和 R0 (静态rho) 的差值 ΔR。")
    
    try:
        # 获取用户输入
        v0 = float(input("\n请输入初速度 v₀ (m/s)："))
        theta = float(input("请输入入射角 θ (度, 0-90)："))
        
        if not (0 < theta <= 90) or v0 <= 0:
            print("输入参数无效。请确保 v₀ > 0 且 0 < θ <= 90。")
            return

    except ValueError:
        print("输入错误。请确保输入的是数字。")
        return
    except Exception as e:
        print(f"发生错误: {e}")
        return

    # --- 计算方法步骤 ---
    print(f"\n--- 📈 计算开始 (dt={dt}s) ---")
    
    # 1. 计算 R1 (动态密度模型)
    R1 = calculate_range(v0, theta, model_type='dynamic', dt=dt)
    print(f"   R1 (考虑 ρ(z) 影响): {R1:.3f} m")

    # 2. 计算 R0 (静态密度模型)
    R0 = calculate_range(v0, theta, model_type='static', dt=dt)
    print(f"   R0 (不考虑 ρ(z) 影响): {R0:.3f} m")

    # 3. 计算偏差 $\Delta R = R1 - R0$
    Delta_R = R1 - R0
    
    print("\n--- 🏁 敏感度分析结果 ---")
    print(f"设定参数: v₀ = {v0:.3f} m/s, θ = {theta:.1f}°")
    print(f"射程偏差 ΔR = R1 - R0 = {Delta_R:.3f} m")
    
    # 偏差分析
    if Delta_R > 0:
        print("分析：由于弹丸在空中大部分时间处于比海平面稀薄的空气中，阻力减小，射程 R1 增加。")
    elif Delta_R < 0:
        print("分析：(通常不发生) 这表明平均密度大于海平面密度，使射程 R1 减小。")
    else:
        print("分析：两种模型射程几乎一致。")

# 启动交互式运行
if __name__ == "__main__":
    interactive_sensitivity_analysis()